/**
 * Copyright 2022 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "src/common/ops/populate/populate_register.h"

#include "nnacl/nllloss_parameter.h"

using mindspore::schema::PrimitiveType_NLLLoss;
using mindspore::schema::PrimitiveType_NLLLossGrad;
using mindspore::schema::Reduction;

namespace mindspore {
namespace lite {
OpParameter *PopulateNLLLossParameter(const void *prim) {
  auto primitive = static_cast<const schema::Primitive *>(prim);
  MS_ASSERT(primitive != nullptr);
  Reduction reduction;
  if (primitive->value_type() == PrimitiveType_NLLLoss) {
    auto value = primitive->value_as_NLLLoss();
    if (value == nullptr) {
      MS_LOG(ERROR) << "value is nullptr";
      return nullptr;
    }
    reduction = value->reduction();
  } else {
    auto value = primitive->value_as_NLLLossGrad();
    if (value == nullptr) {
      MS_LOG(ERROR) << "value is nullptr";
      return nullptr;
    }
    reduction = value->reduction();
  }

  auto *param = reinterpret_cast<NLLLossParameter *>(malloc(sizeof(NLLLossParameter)));
  if (param == nullptr) {
    MS_LOG(ERROR) << "malloc NLLLossParameter failed.";
    return nullptr;
  }
  memset(param, 0, sizeof(NLLLossParameter));

  param->op_parameter_.type_ = primitive->value_type();
  param->reduction_type_ = static_cast<ReductionType>(reduction);
  return reinterpret_cast<OpParameter *>(param);
}

REG_POPULATE(PrimitiveType_NLLLoss, PopulateNLLLossParameter, SCHEMA_CUR)
REG_POPULATE(PrimitiveType_NLLLossGrad, PopulateNLLLossParameter, SCHEMA_CUR)
}  // namespace lite
}  // namespace mindspore
